import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

def load_data():
    fashion_mnist = tf.keras.datasets.fashion_mnist
    (train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()
    # print((train_images, train_labels), (test_images, test_labels))

    return (train_images, train_labels), (test_images, test_labels)



if __name__ == '__main__':

    (train_images, train_labels), (test_images, test_labels) = load_data()

    model = tf.keras.Sequential([
        tf.keras.layers.Flatten(input_shape=(28, 28)),
        tf.keras.layers.Dense(units=128, activation=tf.nn.relu),
        tf.keras.layers.Dense(units=10, activation=tf.nn.softmax)
    ])

    # show img ----------------------
    # img_num = 45
    # print(train_labels[img_num])
    # print(train_images[img_num])
    # plt.imshow(train_images[img_num])
    # plt.show()

    # img_num = 45
    # print(test_labels[img_num])
    # print(test_images[img_num])
    # plt.imshow(test_images[img_num])
    # plt.show()
    # -------------------------------

    model.compile(optimizer='Adam', loss=tf.keras.losses.sparse_categorical_crossentropy)

    model.fit(x=train_images, y=train_labels, epochs=5)

    # res = model.predict(np.reshape(test_images[0], [1, 28, 28]))
    res = model.predict(test_images[:1])

    print(res)
